##########################################################################################

library(data.table)
library(optparse)
library(ArchR)
library(dplyr)

##########################################################################################
option_list <- list(
    make_option(c("--postive_tf_file"), type = "character"),
    make_option(c("--scriptPath"), type = "character"),
    make_option(c("--input_path"), type = "character"),
    make_option(c("--out_path"), type = "character") 
)

if(1!=1){
    
    ## 阳性tf所在文件
    postive_tf_file <- "~/20231121_singleMuti/results/tf_regulators/Positve_TF-Gene.onlyTargetGene.annoExpCor.tsv"

    ## 既往研究整理的代码
    scriptPath <- "~/20231121_singleMuti/scripts/scScalpChromatin"

    ## 每个细胞的Rdata
    input_path <- "~/20231121_singleMuti/results/subcell"

    ## 输出
    out_path <- "~/20231121_singleMuti/results/tf_regulators"

}

###########################################################################################
parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

postive_tf_file <- opt$postive_tf_file
scriptPath <- opt$scriptPath
input_path <- opt$input_path
out_path <- opt$out_path

dir.create(out_path , recursive = T)

###########################################################################################
## 导入数据
postive_tf <- fread(postive_tf_file)

## 已发表文献写好的脚本
source(paste0(scriptPath, "/plotting_config.R"))
source(paste0(scriptPath, "/misc_helpers.R"))
source(paste0(scriptPath, "/matrix_helpers.R"))
source(paste0(scriptPath, "/archr_helpers.R"))
source(paste0(scriptPath, "/GO_wrappers.R"))

###########################################################################################
## 每个clus分别提取TF

result_peak_gene <- c()
result_motifInPeak <- c()

for( clu in unique(postive_tf$cluster) ){
    print(clu)

    ## 该细胞类型中的关键motif
    tmp_motif <- unique(subset( postive_tf , cluster == clu )$motifName)

    ## atac_file
    atac_file <- paste0( input_path , "/" , clu , "/" , clu , ".combineRNA.motif_peak2gene.Rdata" )
    load(atac_file)

    cell_type <- unique(atac_proj@cellColData$cell_type)

    ## 该类细胞中的所有motif
    motifPositions <- getPositions(atac_proj, name="Motif")
    motifGR <- stack(motifPositions, index.var="motifName")

    # Get peak to gene GR
    corrCutoff <- 0.45 # Used in labeling peak2gene links
    p2gGR <- getP2G_GR(atac_proj, corrCutoff=corrCutoff)
        
    for(motif in tmp_motif){
        motif_short <- strsplit(motif,"_")[[1]][1]
        # First get motif positions
        motifLocs <- motifGR[motifGR$motifName == motif]

        # 寻找落在peak上的motif
        ol <- findOverlaps(motifLocs, p2gGR, maxgap=0, type=c("any"), ignore.strand=TRUE)
        olGenes <- p2gGR[to(ol)]
        olGenes$motifScore <- motifLocs[from(ol)]$score
        olGenes$R2 <- olGenes$Correlation**2 
        olGenes$motif <- motif
        olGenes$cluster <- clu
        olGenes$cell_type <- cell_type

        ## 提取存在motif的peak-gene
        olGenes <- data.frame(olGenes)
        result_peak_gene <- bind_rows( result_peak_gene , olGenes )

        ## 提取落在peak-gene上的motif
        motif_in_peak <- data.frame(motifLocs[from(ol)])
        motif_in_peak$cluster <- clu
        motif_in_peak$cell_type <- cell_type

        result_motifInPeak <- bind_rows( result_motifInPeak , motif_in_peak )
    }
}


out_file <- paste0(out_path, "/Positve_TF-Gene.onlyTargetGene.peak-gene_withMotif.tsv")
write.table(result_peak_gene , out_file , row.names = F , sep = "\t" , quote = F)

out_file <- paste0(out_path, "/Positve_TF-Gene.onlyTargetGene.MotifInPeak.tsv")
write.table(result_motifInPeak , out_file , row.names = F , sep = "\t" , quote = F)